/*
Flash.cpp
	
	Controls a MT29F4G01AAADD NAND Flash Device

	Written By Jacob Weiss for Dartmouth Greencube 2013
*/

#include "Flash.h"
#include <Serial.h>
#include <avr/io.h>
#include <EEPROM.h>

#define DEBUG
		
// The location of the page address in EEPROM
#define EEPROM_PAGE_ADDRESS 42
// The location of the byte address in EEPROM
#define EEPROM_BYTE_ADDRESS 36
		
#define FLASH_VALID_ID 11314
		
#define FLASH_START_ADDRESS 64
#define BLOCK_SIZE 64
#define PAGE_SIZE 2048
		
// Pin Numbers / ///////////////////////////////////////////////////////////////

#define SLAVE_SELECT_PIN 8
#define WRITE_PROTECT_PIN 9
#define HOLD_PIN 10
#define MOSI_PIN 11
#define MISO_PIN 12
#define SCK_PIN 13

// Constructors ////////////////////////////////////////////////////////////////

Flash::Flash(bool enable_write, bool restart_address_counter)
{
	// Set up SPI
  SPI.begin(); 
  SPI.setDataMode(SPI_MODE0);
	SPI.setClockDivider(SPI_CLOCK_DIV2);
	
	// Eventually, these can be parameters for the constructor
	this->slaveSelectPin = SLAVE_SELECT_PIN;
	this->writeProtectPin = WRITE_PROTECT_PIN;
	this->holdPin = HOLD_PIN;
	
	// Initialize state constants
	this->write_enabled = enable_write;
	this->dumps_ascii = true;
	
  // Set up output pins and directions
  pinMode(slaveSelectPin, OUTPUT);
	this->csh();
  pinMode(holdPin, OUTPUT);
  digitalWrite(holdPin, HIGH);
  pinMode(writeProtectPin, OUTPUT);
  digitalWrite(writeProtectPin, HIGH);
    
	// Initialize the state
  this->state = STATE_NONE;
	
	if (this->readID() != FLASH_VALID_ID)
	{
		Serial.println("NAND FLASH INITIALIZATION FAILED");
		return;
	}
	
	// Unlock the entire chip
  this->setLocks(LOCK_NONE);		
		
	/*
	If RESTART_ADDRESS_COUNTER is defined, then restart the page address
	to the defined start value and erase the first block. If it is not
	defined, then then continue writing to the flash from the last known
	known location.	
	*/
			
	if (restart_address_counter)
	{		
		this->restartAddressCounter();
				
		if (this->write_enabled)
		{			
			this->blockErase(FLASH_START_ADDRESS);			
			this->blockUntilFinished();			
		}
	}
			
	this->currentPage = this->getEepromPageAddress();			
	this->currentByte = 0;
		
	Serial.print("Initialized to page:");
	Serial.print(this->currentPage);
	Serial.print(" byte:");
	Serial.println(this->currentByte);
}

// Public Methods //////////////////////////////////////////////////////////////


// Write length bytes into the cache. 
// Consecutive calls to writeBytes will continue to fill the cache and will
// not overwrite prevsiously written data. Only 2112
// bytes will fit in the cache. If more than that amount
// of data is written before cacheToArray is called, then
// cacheToArray will be called automatically.
void Flash::writeBytes(byte* data, int length)
{	
	if (this->write_enabled)
	{	
		// If the the state is not program load, then
		// block until the previous operation is completed
		// and start a program load operation.
		if (this->state != STATE_PROGRAM_LOAD)
		{
			this->blockUntilFinished();
			this->programLoadStart(this->getCurrentPlane(), this->currentByte);
		}
		
		// Load the data into the cache register on the NAND device
		int bytesSent;
		for (bytesSent = 0; (bytesSent < length) && (this->currentByte <= PAGE_SIZE - 1); bytesSent++)
		{
			this->programLoad(data[bytesSent]);
			this->currentByte++;
		}
		
		int bytesLeft = length - bytesSent;
		
		// If not all of the bytes were transferred
		if (bytesLeft > 0)
		{
			this->cacheToArray();
			this->blockUntilFinished();			
			this->writeBytes(data + bytesSent, length - bytesSent);
		}
	}
}

// Move the current contents of the cache register into the next available
// page in the main memory array. If the next page is in a new block, then
// the block will be erased prior to writing.
void Flash::cacheToArray()
{	
	if (this->write_enabled)
	{
		// If we are not currently loading a program, then something is very wrong.
		if (this->state == STATE_PROGRAM_LOAD)
		{
			// End the program load operation
			this->programLoadEnd();	
			// Begin a program execute operation
			this->programExecute(this->currentPage++);
		
			// If the next page is the start of a new Block
			// then we will erase it
			if (this->currentPage % BLOCK_SIZE == 0)
			{
				// Wait until the write is finished before trying to erase
				this->blockUntilFinished();
				// Erase the next page (it was incremented when programExecute was called)
				this->blockErase(this->currentPage);
				// Wait until the erase has finished
				this->blockUntilFinished();
			}
				
			// Update the page address stored in the eeprom to reflect the current page.
			this->updateEepromPageAddress(this->currentPage);
			this->currentByte = 0;
			this->updateEepromByteAddress(this->currentByte);
		}
		else
		{
			#ifdef DEBUG
				Serial.println("cacheToArray: State Error");
			#endif
		}
	}
}

// Dumps the array from startAddress up to, but not including, endAddress
// From each page, amount bytes will be dumped. If amount is greater than 2112,
// only 2112 bytes will be dumped.
void Flash::dumpArray(uint32_t startAddress, uint32_t endAddress, int amount, bool dumpAscii)
{
	if (this->changeState(STATE_NONE))
	{
		for (uint32_t i = startAddress; i < endAddress; i++)
		{
			this->blockUntilFinished();
			this->pageRead(i);
			this->blockUntilFinished();
			this->readFromCacheStart(this->getPlane(i), 0);
			for (int j = 0; j < amount && j < PAGE_SIZE; j++)
			{
				if (dumpAscii)
				{
				  Serial.print(this->readFromCache());
				  Serial.print(" ");
				}
				else
				{
					Serial.write(this->readFromCache());
				}
			}
			if (dumpAscii)
			{
			  Serial.println();
			}
			this->readFromCacheEnd();
		}
	}
	else
	{
		#ifdef DEBUG
			Serial.println("dumpArray: State Error");
		#endif
	}
}

void Flash::dumpArray(uint32_t startAddress, uint32_t endAddress, int amount)
{
	this->dumpArray(startAddress, endAddress, amount, false);
}

// Print the id of the nand flash device. It should be 11341.
uint16_t Flash::readID()
{
  this->csl();
  SPI.transfer(READ_ID);
  SPI.transfer(DUMMY_BYTE);
  uint16_t id = SPI.transfer(DUMMY_BYTE) << 8;
  id |= SPI.transfer(DUMMY_BYTE);
  this->csh();
  return id;
}

void Flash::restartAddressCounter()
{
	this->updateEepromPageAddress(FLASH_START_ADDRESS);
	this->updateEepromByteAddress(0);
}

// Private Methods ////////////////////////////////////////////////////////////

// Wait until the current operation is finished. If the operation fails, this will
// loop indefinately.
void Flash::blockUntilFinished()
{
	while(!this->changeState(STATE_NONE));
}

// Write a 32 bit integer to an address in eeprom
void Flash::writeEeprom32(uint32_t data, uint32_t address)
{
	EEPROM.write(address, data & 0xFF);
	EEPROM.write(address + 1, (data >> 8) & 0xFF);
	EEPROM.write(address + 2, (data >> 16) & 0xFF);
	EEPROM.write(address + 3, (data >> 24) & 0xFF);
}

// Read a 32 bit integer from an address in eeprom
uint32_t Flash::readEeprom32(uint32_t address)
{
	return (EEPROM.read(address + 0) << 0)  | 
			   (EEPROM.read(address + 1) << 8)  | 
				 (EEPROM.read(address + 2) << 16) | 
				 (EEPROM.read(address + 3) << 24);
}

// Update the eeprom with a 32 bit page address.
void Flash::updateEepromPageAddress(uint32_t address)
{
	this->writeEeprom32(address, EEPROM_PAGE_ADDRESS);
}

// Get a 32 bit page address from the eeprom.
uint32_t Flash::getEepromPageAddress()
{
	return this->readEeprom32(EEPROM_PAGE_ADDRESS);
}

// Update the eeprom byte counter
void Flash::updateEepromByteAddress(uint32_t address)
{
	this->writeEeprom32(address, EEPROM_BYTE_ADDRESS);
}

// Get a 32 bit byte count from the eeprom.
uint32_t Flash::getEepromByteAddress()
{
	return this->readEeprom32(EEPROM_BYTE_ADDRESS);
}

// Change the state of the flash device if possible. 
// True is returned if the change was successfully made.
// False is returned otherwise.
boolean Flash::changeState(State changeTo)
{
  switch (this->state)
  {
		// If the device is currently performing an operation.
		// Note, block erase, page read, and program execute are
		// the three operations that happen asynchronously.
		case STATE_BLOCK_ERASE:
    case STATE_PAGE_READ:
    case STATE_PROGRAM_EXECUTE:
    {
			// Then check the status of this operation.
			switch (this->operationStatus())
			{
				case OS_IN_PROGRESS:
				{
	        return false;
				}
				case OS_WRITE_FAIL:
				{
					// Serial.println("Write Failure");
					return false;
				}
				case OS_ERASE_FAIL:
				{
					// Serial.println("Erase Failure");
					return false;
				}
				// If the operation was successful
				case OS_SUCCESS:
				{
					// Fall through to STATE_NONE and change the state
					break;
				}
			}
			// Notice there is no break statement here. This is purposeful.
			// The switch falls through to the STATE_NONE case
    }
		// If the device is not performing an operation
    case STATE_NONE:
    {
			// Change the state and return TRUE
      this->state = changeTo;
      return true;
    }
    default:
    {
      return false;
    }
  }
}

// Get the plane of the device being worked on given a page address.
int Flash::getPlane(uint32_t address)
{
	return (address >> 6) & 1;
}
int Flash::getCurrentPlane()
{
	return this->getPlane(this->currentPage);
}

// Set the lock bits of the device
void Flash::setLocks(BlockLock data)
{
	this->csl();
  SPI.transfer(SET_FEATURE);
  SPI.transfer(FEATURE_BLOCK_LOCK);
  SPI.transfer(data & 0xFF);
  this->csh();
}

// Enable write/erase operations
void Flash::writeEnable()
{
  this->csl();
  SPI.transfer(WRITE_ENABLE);
  this->csh();
}

// Send a program execute
void Flash::programExecute(uint32_t address)
{
  if (this->changeState(STATE_PROGRAM_EXECUTE))
  {
    this->csl();
    SPI.transfer(PROGRAM_EXECUTE);
    SPI.transfer((address >> 16) & 0x03);
    SPI.transfer((address >> 8) & 0xFF);
    SPI.transfer(address & 0xFF);
    this->csh();
  }
  else
  {
		#ifdef DEBUG
			Serial.println("programExecute: State Failure");
		#endif
  }
}

// Begin loading bytes into the cache in the given plane
// and column
void Flash::programLoadStart(int plane, uint16_t column)
{
  if (this->changeState(STATE_PROGRAM_LOAD))
  {
    this->writeEnable();
    uint16_t address = ((plane & 0x01) << 12) | (column & 0x0FFF);
    this->csl();
    SPI.transfer(PROGRAM_LOAD);
    SPI.transfer((address >> 8) & 0xFF); // Shift out 8 high bits
    SPI.transfer(address && 0xFF); // Shift out 8 low bits
    this->state = STATE_PROGRAM_LOAD;
  }
  else
  {
		#ifdef DEBUG
			Serial.println("programLoadStart: State Failure");
		#endif
  }
}

// Load a single byte into the cache.
void Flash::programLoad(byte data)
{
  if (this->state == STATE_PROGRAM_LOAD)
  {
    SPI.transfer(data);
  }
  else
  {
		#ifdef DEBUG
			Serial.println("programLoad: State Failure");
		#endif
  }
}

// Finish loading bytes into the cache.
void Flash::programLoadEnd()
{
  if (this->state == STATE_PROGRAM_LOAD)
  {
    this->csh();
    state = STATE_NONE;
  }
  else
  {
		#ifdef DEBUG
				Serial.println("programLoadEnd: State Failure");
		#endif
  }
}

// Begin reading bytes from the cache at the given plane and column
void Flash::readFromCacheStart(int plane, uint16_t column)
{
  if (this->changeState(STATE_READ_FROM_CACHE))
  {
    uint16_t address = ((plane & 0x01) << 12) | (column & 0x0FFF); 
    this->csl();
    SPI.transfer(READ_FROM_CACHE);
    SPI.transfer((address >> 8) & 0xFF); // Shift out 8 high bits
    SPI.transfer(address && 0xFF); // Shift out 8 low bits
    SPI.transfer(DUMMY_BYTE);   
  }
  else
  {
		#ifdef DEBUG
			Serial.println("readFromCacheStart: State Failure");
		#endif
  }
}

// Read a byte from the cache
int Flash::readFromCache()
{
  if (this->state == STATE_READ_FROM_CACHE)
  {
    return SPI.transfer(DUMMY_BYTE);
  }
  else
  {
		#ifdef DEBUG
			Serial.println("readFromCache: State Failure");
		#endif
  }
  return 0;
}

// Finish reading bytes from the cache
void Flash::readFromCacheEnd()
{
  if (this->state == STATE_READ_FROM_CACHE)
  {
    this->csh();
    state = STATE_NONE;
  }
  else
  {
		#ifdef DEBUG
   		Serial.println("readFromCacheEnd: State Failure");
		#endif
  }
}

// Load a page from the main array into the cache
void Flash::pageRead(uint32_t address)
{
  if (this->changeState(STATE_PAGE_READ))
  {
    this->csl();
    SPI.transfer(PAGE_READ);
    SPI.transfer((address >> 16) & (0x03)); // Only the most significant 2 bits are valid.
    SPI.transfer((address >> 8) & (0xFF)); 
    SPI.transfer((address) & (0xFF));
    this->csh();
  }
  else
  {
		#ifdef DEBUG
    	Serial.println("pageRead: State Failure");
		#endif	
  }
}

// Erase a single block in the main array
void Flash::blockErase(uint32_t address)
{
	if (this->changeState(STATE_BLOCK_ERASE))
	{
		this->writeEnable();
		this->csl();
		SPI.transfer(BLOCK_ERASE);
    SPI.transfer((address >> 16) & (0x03)); // Only the most significant 2 bits are valid.
    SPI.transfer((address >> 8) & (0xFF)); 
    SPI.transfer((address) & (0xFF));
		this->csh();
	}
  else
  {
		#ifdef DEBUG
			Serial.println("blockErase: State Failure");
		#endif
  }
}

// Reset the flash device
void Flash::reset()
{
  this->csl();
  SPI.transfer(RESET);
  this->csh();
}

// Get the current status of an operation
OperationStatus Flash::operationStatus()
{
  int f = getFeature(FEATURE_STATUS); //& 0x01;
	if (f & 0x01)
	{
		return OS_IN_PROGRESS;
	}
	else if (f & 0x04)
	{
		return OS_ERASE_FAIL;
	}
	else if (f & 0x08)
	{
		return OS_WRITE_FAIL;
	}
  return OS_SUCCESS;
}

// Read a feature byte
int Flash::getFeature(FeatureByte address)
{
  this->csl();
  SPI.transfer(GET_FEATURE);
  SPI.transfer(address & 0xFF);
  int feature = SPI.transfer(DUMMY_BYTE);
  this->csh();
  return feature;
}

// Bring the chip select pin high
void Flash::csh()
{
  digitalWrite(this->slaveSelectPin, HIGH);
}
/// Bring the chip select pin low
void Flash::csl()
{
  digitalWrite(this->slaveSelectPin, LOW);
}
